import os
import base64
from chromadb import PersistentClient
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
from typing import List, Dict, Any
from PIL import Image
import io

# 1. 初始化Chroma客户端和多模态嵌入函数
class MultiModalEmbeddingFunction(EmbeddingFunction):
    def __init__(self):
        #self.clip_ef = OpenCLIPEmbeddingFunction()
        self.clip_ef = SentenceTransformerEmbeddingFunction(model_name="D:/ideaSpace/MyPython/models/all-MiniLM-L6-v2")

    def __call__(self, input: Documents) -> Embeddings:
        # 处理混合类型输入（文本或base64编码图像）
        embeddings = []
        for item in input:
            if item.startswith("data:image"):  # 假设是base64编码图像
                embeddings.append(self.clip_ef([item])[0])
            else:  # 文本
                embeddings.append(self.clip_ef([item])[0])
        return embeddings

# 连接到Chroma
client = PersistentClient(path="./chroma_multimodal_db")

# 检查并创建集合
collection_name = "Monkey"
if collection_name in [col.name for col in client.list_collections()]:
    client.delete_collection(collection_name)

collection = client.create_collection(
    name=collection_name,
    embedding_function=MultiModalEmbeddingFunction()
)

# 2. 数据处理函数
def to_base64(path: str) -> str:
    """将文件转换为base64编码字符串"""
    with open(path, "rb") as f:
        return f"data:image/jpeg;base64,{base64.b64encode(f.read()).decode('utf-8')}"

def load_image_data(image_dir: str) -> List[Dict[str, Any]]:
    """加载图像数据并生成元数据"""
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    data = []
    for name in image_files:
        path = os.path.join(image_dir, name)
        data.append({
            "id": name,
            "document": to_base64(path),  # 使用base64编码作为文档
            "metadata": {
                "path": path,
                "mediaType": "image",
                "name": name
            }
        })
    return data

# 3. 插入图像数据
image_dir = "D:/ideaSpace/rag-in-action-master/90-文档-Data/多模态/Weaviate"  # 替换为你的实际路径
image_data = load_image_data(image_dir)

# 批量插入数据
collection.add(
    ids=[item["id"] for item in image_data],
    documents=[item["document"] for item in image_data],
    metadatas=[item["metadata"] for item in image_data]
)

# 4. 查询功能
def text_search(query: str, limit: int = 3):
    """文本搜索相似图像"""
    results = collection.query(
        query_texts=[query],
        n_results=limit,
        include=["metadatas", "distances"]
    )
    print(f"\n文本查询: '{query}' 的相似结果:")
    for meta in results["metadatas"][0]:
        print(f"- {meta['name']} (路径: {meta['path']})")

def image_search(image_path: str, limit: int = 3):
    """图像搜索相似图像"""
    query_image = to_base64(image_path)
    results = collection.query(
        query_texts=[query_image],  # 使用相同的嵌入函数处理图像
        n_results=limit,
        include=["metadatas", "distances"]
    )
    print(f"\n图像查询: '{image_path}' 的相似结果:")
    for meta in results["metadatas"][0]:
        print(f"- {meta['name']} (路径: {meta['path']})")

# 5. 执行查询示例
# 文本查询
text_search("Monkey with fire")
text_search("Monsters")

# 图像查询
test_image_path = "D:/ideaSpace/rag-in-action-master/90-文档-Data/多模态/query_image.jpg"  # 替换为你的查询图像路径
if os.path.exists(test_image_path):
    image_search(test_image_path)
else:
    print(f"\n测试图像不存在: {test_image_path}")

# 6. 显示结果函数（可选）
def display_results(results):
    """显示查询结果图像"""
    try:
        from IPython.display import display, Image
        for meta in results["metadatas"][0]:
            display(Image(filename=meta["path"]))
    except ImportError:
        print("IPython未安装，无法显示图像")

# 关闭连接（Chroma的PersistentClient不需要显式关闭）